import os
import librosa
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, confusion_matrix
import seaborn as sns
import matplotlib.pyplot as plt
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Dropout
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.callbacks import ModelCheckpoint

# 设置输出文件夹
output_dir = "output_2"
os.makedirs(output_dir, exist_ok=True)

# 数据集路径
dataset_path = "UrbanSound8K"
metadata_path = os.path.join(dataset_path, "metadata/UrbanSound8K.csv")
audio_path = os.path.join(dataset_path, "audio")

# 加载元数据
metadata = pd.read_csv(metadata_path)

# 限制每个类别的样本数量（每类 100 条）
metadata = metadata.groupby("classID").head(100)


# 提取音频特征（MFCC）
def extract_features(file_path):
    y, sr = librosa.load(file_path, sr=None, mono=True)
    mfccs = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=40)
    return np.mean(mfccs.T, axis=0)


# 提取所有音频文件的特征
features = []
labels = []

for index, row in metadata.iterrows():
    file_path = os.path.join(audio_path, f"fold{row['fold']}", row['slice_file_name'])
    try:
        mfcc = extract_features(file_path)
        features.append(mfcc)
        labels.append(row["classID"])
    except Exception as e:
        print(f"Error processing {file_path}: {e}")

# 转换为 NumPy 数组
X = np.array(features)
y = np.array(labels)

# 将数据划分为训练集、验证集和测试集
X_train, X_temp, y_train, y_temp = train_test_split(X, y, test_size=0.3, random_state=42)  # 70% 训练集
X_val, X_test, y_val, y_test = train_test_split(X_temp, y_temp, test_size=0.5, random_state=42)  # 15% 验证集，15% 测试集

# 将标签转换为独热编码
y_train = to_categorical(y_train)
y_val = to_categorical(y_val)
y_test = to_categorical(y_test)

# 构建神经网络模型
model = Sequential([
    Dense(256, activation='relu', input_shape=(X_train.shape[1],)),
    Dropout(0.3),
    Dense(128, activation='relu'),
    Dropout(0.3),
    Dense(len(np.unique(labels)), activation='softmax')  # 输出类别数为所有类别数
])

# 编译模型
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

# 设置模型保存路径
model_path = os.path.join(output_dir, "best_model.h5")
checkpoint = ModelCheckpoint(model_path, monitor='val_loss', save_best_only=True, verbose=1)

# 训练模型
history = model.fit(X_train, y_train, epochs=30, batch_size=32, validation_data=(X_val, y_val), callbacks=[checkpoint])

# 测试模型
test_loss, test_accuracy = model.evaluate(X_test, y_test)
print(f"Test Accuracy: {test_accuracy * 100:.2f}%")

# 获取预测结果
y_pred = model.predict(X_test)
y_pred_classes = np.argmax(y_pred, axis=1)  # 将预测结果转换为类别索引
y_true = np.argmax(y_test, axis=1)  # 将真实标签转换为类别索引

# 打印分类报告并保存到文件
classification_report_str = classification_report(y_true, y_pred_classes)
print("\nClassification Report:")
print(classification_report_str)
report_path = os.path.join(output_dir, "classification_report.txt")
with open(report_path, "w") as f:
    f.write(classification_report_str)

# 绘制混淆矩阵并保存到文件
conf_matrix = confusion_matrix(y_true, y_pred_classes)
plt.figure(figsize=(10, 8))
sns.heatmap(conf_matrix, annot=True, fmt='d', cmap='Blues', xticklabels=np.unique(labels), yticklabels=np.unique(labels))
plt.xlabel('Predicted Labels')
plt.ylabel('True Labels')
plt.title('Confusion Matrix')
conf_matrix_path = os.path.join(output_dir, "confusion_matrix.png")
plt.savefig(conf_matrix_path)
plt.close()

# 保存模型参数
model_summary_path = os.path.join(output_dir, "model_summary.txt")
with open(model_summary_path, "w") as f:
    model.summary(print_fn=lambda x: f.write(x + "\n"))

print(f"所有结果已保存到文件夹: {output_dir}")